from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.asymmetric import padding
import base64

""" 
生成rsa 公钥,密钥
openssl genpkey -algorithm RSA -out private_key.pem -pkeyopt rsa_keygen_bits:2048
openssl rsa -in private_key.pem -pubout -outform PEM -out public_key.pem
"""

def load_private_key(filename):
    """从文件加载私钥"""
    with open(filename, "rb") as key_file:
        private_key = serialization.load_pem_private_key(
            key_file.read(),
            password=None,  # 如果你的私钥没有密码，这里设置为None
        )
    return private_key


def load_public_key(filename):
    """从文件加载公钥"""
    with open(filename, "rb") as key_file:
        public_key = serialization.load_pem_public_key(key_file.read())
    return public_key


def encrypt_message(message, public_key):
    """使用公钥加密消息"""
    ciphertext = public_key.encrypt(
        message.encode("utf-8"),
        padding.OAEP(
            mgf=padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None,
        ),
    )
    return base64.b64encode(ciphertext).decode("utf-8")


def decrypt_message(encrypted_message, private_key):
    """使用私钥解密消息"""
    ciphertext = base64.b64decode(encrypted_message)
    plaintext = private_key.decrypt(
        ciphertext,
        padding.OAEP(
            mgf=padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None,
        ),
    )
    return plaintext.decode("utf-8")

buf = "hello world"
public_key = load_public_key("public_key.pem")
private_key = load_private_key("private_key.pem")
encode_buf = encrypt_message(buf,public_key)
print("encode_buf len=",len(encode_buf))

decode_buf = decrypt_message(encode_buf,private_key)
print("decode_buf")
print(decode_buf)